package com.example.hyvelucenedemo.utils;

import com.example.hyvelucenedemo.pojo.IndexItem;
import com.example.hyvelucenedemo.pojo.SearchItem;
import com.example.hyvelucenedemo.pojo.SearchPageVO;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.*;
import org.apache.lucene.index.*;
import org.apache.lucene.queryparser.classic.MultiFieldQueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.FSDirectory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.wltea.analyzer.lucene.IKAnalyzer;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @Description Lucene搜索工具类
 * @Duthor jie.zhao
 * @Date 2021/5/9 11:07 上午
 */
@Component
public class LuceneSearchUtil implements InitializingBean {

    @Value("${index.path}")
    private String indexPath;

    /**
     * 搜索
     *
     * @param keyword    关键字
     * @param pageNumber 当前页
     * @param pageSize   每页的数量
     * @return
     */
    public SearchPageVO searchPageVO(String keyword, int pageNumber, int pageSize) {
        SearchPageVO searchPageVO = new SearchPageVO(pageNumber, pageSize);
        IndexReader indexReader = createIndexRead();
        try {
            //创建IndexSearcher
            IndexSearcher indexSearcher = new IndexSearcher(indexReader);
            Query query = MultiFieldQueryParser.parse(new String[]{keyword, keyword, keyword}, new String[]{"title", "content", "labels"}, new IKAnalyzer());
            QueryScorer scorer = new QueryScorer(query, null);
            SimpleHTMLFormatter simpleHTMLFormatter = new SimpleHTMLFormatter("<span class='red'>", "</span>");
            Highlighter highlighter = new Highlighter(simpleHTMLFormatter, scorer);
            Fragmenter fragment = new SimpleSpanFragmenter(scorer);
            highlighter.setTextFragmenter(fragment);
            //分页查询
            if (pageNumber < 1) {
                pageNumber = 1;
            }
            if (pageSize < 1) {
                pageSize = 10;
            }
            TopDocs topDocs;
            int start = (pageNumber - 1) * pageSize;
            if (0 == start) {
                topDocs = indexSearcher.search(query, pageSize);
            } else {
                //查询数据，结束页面自前的数据都会查询到，但是只取本页的数据
                topDocs = indexSearcher.search(query, start);
                //获取到上一页最后一条
                ScoreDoc preScore = topDocs.scoreDocs[start - 1];
                //查询最后一条后的数据的一页数据
                topDocs = indexSearcher.searchAfter(preScore, query, pageSize);
            }
            //查询结果
            List<SearchItem> result = new ArrayList<>();
            for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
                Document document = indexSearcher.doc(scoreDoc.doc);
                result.add(convertDocument(document, highlighter));
            }
            //总条数
            long totalElements = topDocs.totalHits.value;
            //总页数
            long totalPages = (totalElements + pageSize - 1) / pageSize;
            //组装返回数据
            searchPageVO.setContent(result);
            searchPageVO.setTotalElements(totalElements);
            searchPageVO.setTotalPages(totalPages);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                if (indexReader != null) {
                    indexReader.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return searchPageVO;
    }

    /**
     * 单个创建或更新索引
     * <p>
     * 性能低适合单个索引的创建和更新，不适合批量添加。
     *
     * @param indexItem 索引Item
     * @param isCover   是否覆盖
     */
    public void buildIndex(IndexItem indexItem, boolean isCover) {
        IndexWriter indexWriter = createIndexWrite();
        try {
            Document document = convertIndexItem(indexItem);
            if (isCover) {
                //根据onlyId 更新索引，不存在添加！！为防止重复索引，需要保证onlyId的全局唯一性
                Term onlyIdTerm = new Term("onlyId", indexItem.getOnlyId());
                indexWriter.updateDocument(onlyIdTerm, document);
            } else {
                //添加索引
                indexWriter.addDocument(document);
            }
            indexWriter.commit();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (indexWriter != null) {
                    indexWriter.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 删除索引
     *
     * @param onlyId 唯一ID
     */
    public void deleteIndex(String onlyId) {
        IndexWriter indexWriter = createIndexWrite();
        IndexReader indexReader = createIndexRead();
        try {
            //删除索引，注意要关闭IndexReader，否则不会立即生效
            Term onlyIdTerm = new Term("onlyId", onlyId);
            indexWriter.deleteDocuments(onlyIdTerm);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (indexWriter != null) {
                    indexWriter.close();
                }
                if (indexReader != null) {
                    indexReader.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }


    /**
     * 批量创建索引，注意无法根据onlyId查重需要自己做处理。
     * <p>
     * 性能高 10w数据 10s左右
     *
     * @param indexItems
     */
    public void batchCreateIndex(List<IndexItem> indexItems) {
        IndexWriter indexWriter = createIndexWrite();
        try {
            List<Document> documents = new ArrayList<>();
            indexItems.forEach(indexItem -> {
                documents.add(convertIndexItem(indexItem));
            });
            //添加索引
            indexWriter.addDocuments(documents);
            indexWriter.commit();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                if (indexWriter != null) {
                    indexWriter.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 创建Document
     * <p>
     * StringField 不分词 索引
     * TextField   分词   索引
     *
     * @param indexItem
     * @return
     */
    public Document convertIndexItem(IndexItem indexItem) {
        FieldType onlyType = new FieldType();
        onlyType.setIndexOptions(IndexOptions.DOCS);
        onlyType.setStored(true);

        Field field1 = new Field("onlyId", indexItem.getOnlyId(), onlyType);
        Field field2 = new StringField("id", indexItem.getId(), Field.Store.YES);
        Field field3 = new StringField("type", indexItem.getType(), Field.Store.YES);
        Field field4 = new TextField("labels", indexItem.getLabels(), Field.Store.YES);
        Field field5 = new TextField("title", indexItem.getTitle(), Field.Store.YES);
        Field field6 = new TextField("content", indexItem.getContent(), Field.Store.YES);

        Document document = new Document();
        document.add(field1);
        document.add(field2);
        document.add(field3);
        document.add(field4);
        document.add(field5);
        document.add(field6);
        return document;
    }

    /**
     * 创建SearchItem
     *
     * @param document
     * @param highlighter
     * @return
     * @throws InvalidTokenOffsetsException
     * @throws IOException
     */
    private SearchItem convertDocument(Document document, Highlighter highlighter) throws InvalidTokenOffsetsException, IOException {
        SearchItem searchItem = new SearchItem();
        searchItem.setId(document.get("id"));
        searchItem.setType(document.get("type"));
        searchItem.setLabels(document.get("labels"));
        searchItem.setTitle(document.get("title"));
        searchItem.setContent(document.get("content"));

        // 获取相对的内容
        String content = document.get("content");
        TokenStream tokenStream = new IKAnalyzer(true).tokenStream("content", content);
        String str[] = highlighter.getBestFragments(tokenStream, document.get("content"), 3);//获取高亮的片段，可以对其数量进行限制
        String describe = Arrays.stream(str).collect(Collectors.joining("</br>"));
        //没有高亮词，说明是通过其他字段进行查询到的
        searchItem.setDescribe(describe);
        return searchItem;
    }

    /**
     * 创建IndexWriter，注意每次使用完要close
     *
     * @return
     */
    private IndexWriter createIndexWrite() {
        FSDirectory fsDirectory;
        try {
            fsDirectory = FSDirectory.open(new File(indexPath).toPath());
            IndexWriterConfig config = new IndexWriterConfig(new IKAnalyzer()); //使用IK分词器
            IndexWriter indexWriter = new IndexWriter(fsDirectory, config);
            return indexWriter;
        } catch (IOException e) {
            throw new RuntimeException("打开lucene索引异常", e);
        }
    }

    /**
     * 创建IndexReader，注意每次使用完要close
     *
     * @return
     */
    private IndexReader createIndexRead() {
        FSDirectory fsDirectory;
        try {
            fsDirectory = FSDirectory.open(new File(indexPath).toPath());
            IndexReader indexReader = DirectoryReader.open(fsDirectory);
            return indexReader;
        } catch (IOException e) {
            throw new RuntimeException("打开lucene索引异常", e);
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        // 初始目录
        File file = new File(indexPath);
        if (!file.exists()) {
            file.mkdirs();
        }
    }
}
